from Abstract import *
from Instructions import *
from Arguments import *


class NFAction (Action):
	Command = 'iptables'
	Target = None
	Children = (Inbound, Outbound, Source, Destination, Protocol, ICMP_Type, State, Comment, Options)

	def __init__ (self, Filename, LineNo):
		super (NFAction, self).__init__ (Filename, LineNo)
		self.Matches = []

	def check (self):
		for AP_Class in (Source, Destination, Rewrite_Spec):
			AP = self.child (AP_Class)
			if AP:
				AP_Ports = AP.child (Ports)
				if AP_Ports:
					self.enforce_proto (('tcp', 'udp', '6', '17'), Ports)
					break
		if self.child (ICMP_Type):
			self.enforce_proto (('icmp', '1'), ICMP_Type)

	def enforce_proto (self, Protocols, ReasonClass):
		tProtocol = self.child (Protocol)
		tOptions = self.child (Options)
		Correct = False
		GuessedProto = None
		if tProtocol and not tProtocol.Invert:
			Correct = True
			for Proto in tProtocol.Values:
				if Proto not in Protocols:
					Correct = False
					GuessedProto = Proto
					break
		if not Correct and tOptions:
			Found = False
			Prev = None
			for Option in tOptions.Values:
				if Found:
					if Option.lower () in Protocols:
						Correct = True
					else:
						GuessedProto = Option
					break
				else:
					if Option.lower () in ('-p', '--protocol'):
						if Prev == '!':
							break
						else:
							Found = True
				Prev = Option
		if not Correct:
			self.fail_check ("Protocol required by '%s' token is one of: %s; found '%s' instead" % (ReasonClass.Token, Protocols, GuessedProto))

	def render (self):
		def add (Ax, Op):
			Result = []
			for Item in Ax:
				Item.extend (Op)
				Result.append (Item)

		def multiply (Ax, Op):
			Result = []
			for b in Op:
				for a in Ax:
					c = a + [b,]
					Result.append (c)
			return (Result)

		Opts = [[self.Command,]]
		myInbound = self.child (Inbound)
		myOutbound = self.child (Outbound)
		myNew = self.child (Rewrite_Spec)
		if myNew:
			add (Opts, ('--table', 'nat'))
			myNewAddr = myNew.child (IPv4)
			myNewPorts = myNew.child (Ports)
			myNewRandom = myNew.child (Random)
			myNewPersist = myNew.child (Persist)
			tMyself = False
			if myInbound:
				add (Opts, ('--append', 'PREROUTING'))
				if myNewAddr and '(myself)' in myNewAddr.Values:
					add (Opts, ('--jump', 'REDIRECT'))
					tMyself = True
				else:
					add (Opts, ('--jump', 'DNAT'))
			else:
				add (Opts, ('--append', 'POSTROUTING'))
				if myNewAddr and '(myself)' in myNewAddr.Values:
					add (Opts, ('--jump', 'MASQUERADE'))
					tMyself = True
				else:
					add (Opts, ('--jump', 'SNAT'))
			if tMyself:
				if myNewPorts:
					add (Opts, (['--to-ports',] + [Port.replace ('..', '-') for Port in myNewPorts.Values]))
			else:
				if myNewAddr:
					tAddr, = myNewAddr.Values
					if '/' in tAddr:
						tAddr = IPv4.net2range (tAddr)
					if myNewPorts:
						tPorts, = myNewPorts.Values
						tAddr = '%s:%s' % (tAddr, tPorts)
					add (Opts, ('--to-destination', tAddr.replace ('..', '-')))
				if myNewPersist:
					add (Opts, ('--persistent',))
			if myNewRandom:
				add (Opts, ('--random',))
		else:
			add (Opts, ('--table', 'filter'))
			if myInbound and myOutbound:
				add (Opts, ('--append', 'FORWARD'))
			elif myInbound:
				add (Opts, ('--append', 'INPUT'))
			elif myOutbound:
				add (Opts, ('--append', 'OUTPUT'))
			else:
				add (Opts, ('--append',))
				Opts = multiply (Opts, ('INPUT', 'FORWARD', 'OUTPUT'))
			add (Opts, ('--jump', self.Target))

		for tDir in (myInbound, myOutbound):
			if tDir:
				tIfaces = tDir.child (Interface)
				if tIfaces:
					if tIfaces.Invert:
						add (Opts, ('!',))
					add (Opts, ('--%s-interface' % tDir.Token,))
					Opts = multiply (Opts, tIfaces.Values)

		mySrc = self.child (Source)
		myDst = self.child (Destination)
#		for tAddrPorts in (mySrc, myDst):
#			if 

		myOptions = self.child (Options)
		if myOptions:
			add (Opts, myOptions.Values)

		myComment = self.child (Comment)
		if myComment:
			add (Opts, ['--comment',] + myComment.Values)

		Lines = []
		for Line in Opts:
			Lines.append (' '.join (Argument.escape (Option) for Option in Line))
		return '\n'.join (Lines)
# ----------------------------------------------------------------------------

class IPTables (NFAction):
	Token = 'iptables'
	Children = (Comment, Options)

	def check (self):
		super (IPTables, self).check ()
		if not self.child (Options):
			self.fail_check ("'%s' action requires '%s' instruction" % (self.Token, Options.Token))

	def render (self):
		Opts = [self.Command,]
		myOptions = self.child (Options)
		myComment = self.child (Comment)
		Opts.extend (myOptions.Values)
		if myComment:
			Opts.extend (['--comment',] + myComment.Values)
		return (' '.join (Argument.escape (Option) for Option in Opts))
# ----------------------------------------------------------------------------

class Pass (NFAction):
	Token = 'pass'
	Target = 'ACCEPT'
# ----------------------------------------------------------------------------

class Drop (NFAction):
	Token = 'drop'
	Target = 'DROP'
# ----------------------------------------------------------------------------

class Reject (NFAction):
	Token = 'reject'
	Target = 'REJECT'
	Children = (Inbound, Outbound, Source, Destination, Protocol, ICMP_Type, State, Comment, Options, Reject_Type)
# ----------------------------------------------------------------------------

class Rewrite (NFAction):
	Token = 'rewrite'
	Children = (Inbound, Outbound, Source, Destination, Protocol, ICMP_Type, State, Comment, Options, Rewrite_Spec)

	def check (self):
		super (Rewrite, self).check ()
		if not self.child (Inbound) and not self.child (Outbound):
			self.fail_check ("'%s' action requires either '%s' or '%s' instruction" % (self.Token, Inbound.Token, Outbound.Token))
		if self.child (Inbound) and self.child (Outbound):
			self.fail_check ("'%s' action requires either '%s' or '%s' instruction, not both" % (self.Token, Inbound.Token, Outbound.Token))
		if not self.child (Rewrite_Spec):
			self.fail_check ("'%s' action requires '%s' instruction" % (self.Token, Rewrite_Spec.Token))
# ----------------------------------------------------------------------------
